from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from .enum_class import EmbeddingType


class EmbeddingHandler:
    """
    嵌入模型处理器
    """

    def __init__(self, embedding_type: EmbeddingType, **kwargs):
        self.embedding_type = embedding_type
        self.embedding_model_name = None
        if "embedding_model_name" in kwargs:
            self.embedding_model_name = kwargs["embedding_model_name"]
            # 在kwargs中删除embedding_model_name
            kwargs.pop("embedding_model_name")
        self.kwargs = kwargs

    def get_embedding(self):
        """
        获取嵌入模型，统一入口
        """
        match self.embedding_type:
            case EmbeddingType.Local_HuggingFace:
                return self.get_huggingface_embedding(**self.kwargs)
            case _:
                raise ValueError(f"不支持的嵌入模型类型: {self.embedding_type}")

    def get_huggingface_embedding(self, **kwargs):
        """
        获取HuggingFace的嵌入模型
        参数：
            kwargs: 其他参数
                可选参数:
                    embedding_model_name: 模型名称，默认"BAAI/bge-small-zh-v1.5"，生产环境建议使用更大的模型
        返回：
            HuggingFaceEmbedding: HuggingFace的Embedding模型
        """
        default_model_name = "BAAI/bge-small-zh-v1.5"
        return HuggingFaceEmbedding(
            model_name=self.embedding_model_name
            if self.embedding_model_name
            else default_model_name,
        )
